import numpy as np

np.random.seed(0)

P = [
    [0.9, 0.1, 0.0, 0.0, 0.0, 0.0],
    [0.5, 0.0, 0.5, 0.0, 0.0, 0.0],
    [0.0, 0.0, 0.0, 0.6, 0.0, 0.4],
    [0.0, 0.0, 0.0, 0.0, 0.3, 0.7],
    [0.0, 0.2, 0.3, 0.5, 0.0, 0.0],
    [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
]

p = np.array(P)

rewards = [-1, -2, -2, 10, 1, 0]
gamma = 0.5  # 折扣因子

def compute_return(start_index, chain, gamma):
    """根据一条马尔可夫链，计算回报"""
    G = 0
    for i in reversed(range(start_index, len(chain))):
        G = gamma * G +  rewards[chain[i] - 1]
    return G

chain = [1, 2, 3, 6] # 一条马尔可夫链s1-s2-s3-s6
start_index = 0
G = compute_return(start_index, chain, gamma)
print("回报:%s" % G)

def compute_v(P, rewards, gamma, states_num):
    """利用贝尔曼方程计算状态值"""
    rewards = np.array(rewards).reshape((-1, 1)) # 转换为列向量
    value = np.dot(np.linalg.inv(np.eye(states_num, states_num) - gamma * P), rewards)
    return value

V = compute_v(p, rewards, gamma, 6)
print("状态值：%s" % V)